"""
根据csv路网切割TIF，并算出路宽，TIF和csv需在同一坐标，可以先配准
"""

import csv
import osgeo.gdal as gdal
import numpy as np
import os
import get_direction
import main_point
import red_point
import inner_circle
import cv2
from tqdm import tqdm


def crop_by_csv(csv_file, tif_file, cropped_dir, output_dir, circle_dir):
    # 读取原始 CSV 文件
    with open(csv_file, 'r', encoding='gbk') as csvfile:
        reader = csv.reader(csvfile)
        rows = [row for row in reader]
        rows[0].append('pixel_num')
        # 遍历每一行
        for row in tqdm(rows[1:], desc='Processing rows'):
            wkt_polyline = row[3]  # 获取 WKT_POLYLINE 列值
            coords = wkt_polyline.replace('LINESTRING (', '').replace(')', '').split(', ')
            coords = [coord.split(' ') for coord in coords]
            coords = [(float(coord[0]), float(coord[1])) for coord in coords]
            point1 = coords[0]
            point2 = coords[-1]
            # print('p1:'+str(point1)+' p2:'+str(point2))
            # 打开图像文件
            ds = gdal.Open(tif_file)

            # 获取地理变换矩阵
            gt = ds.GetGeoTransform()

            # 获取逆地理变换矩阵
            inv_gt = gdal.InvGeoTransform(gt)

            # 将坐标转换为像素坐标
            ulx, uly = gdal.ApplyGeoTransform(inv_gt, point1[0], point1[1])
            lrx, lry = gdal.ApplyGeoTransform(inv_gt, point2[0], point2[1])

            # 计算裁剪范围
            x_min, x_max = min(ulx, lrx), max(ulx, lrx)
            y_min, y_max = min(uly, lry), max(uly, lry)
            # print('x: '+str(x_min),str(x_max))
            # print('y: '+str(y_min),str(y_max))
            # 读取图像数据
            band = ds.GetRasterBand(1)
            data = band.ReadAsArray()

            # 扩充范围
            if x_max - x_min > y_max - y_min:
                y_max = min(y_max + 20, data.shape[0])
                y_min = max(y_min - 20, 0)
            else:
                x_max = min(x_max + 40, data.shape[1])
                x_min = max(x_min - 20, 0)

            if y_max - y_min < 18:
                y_max = min(y_max + 10, data.shape[0])
                y_min = max(y_min - 10, 0)
            if x_max - x_min < 18:
                x_max = min(x_max + 10, data.shape[1])
                x_min = max(x_min - 10, 0)

            # 裁剪图像
            cropped_data = data[int(y_min):int(y_max), int(x_min):int(x_max)]

            # 计算白色像素的个数
            white_pixels = np.sum(cropped_data == 255)

            # if white_pixels < 6:
            #     print(f"跳过行 {row[0]}，因为白色像素个数小于 6")
            #     continue

            # 检查cropped_data是否为空
            if cropped_data.size == 0:
                print(f"跳过行 {row[0]}，因为cropped_data为空")
                continue
            # 创建新的图像文件
            driver = gdal.GetDriverByName('GTiff')
            new_ds = driver.Create(f'{cropped_dir}/cropped_image_{row[0]}.tif', cropped_data.shape[1],
                                   cropped_data.shape[0], 1, gdal.GDT_Byte)

            # 将裁剪数据写入新的图像文件
            new_band = new_ds.GetRasterBand(1)
            new_band.WriteArray(cropped_data)

            # 更新地理变换矩阵
            new_gt = list(gt)
            new_gt[0] = min(point1[0], point2[0])
            new_gt[3] = max(point1[1], point2[1])
            new_ds.SetGeoTransform(new_gt)

            # 设置投影信息
            new_ds.SetProjection(ds.GetProjection())
            new_ds.FlushCache()
            new_ds = None

    #         # 读取图像
    #         image = cv2.imread(f'{cropped_dir}/cropped_image_{row[0]}.tif', 0)  # 以灰度模式读取图像
    #
    #         # 定义膨胀核（结构元素）
    #         kernel = np.ones((4, 4), np.uint8)
    #
    #         # 应用膨胀操作
    #         dilation = cv2.dilate(image, kernel, iterations=1)
    #
    #         cv2.imwrite(f'{cropped_dir}_dila/cropped_image_{row[0]}.tif', dilation)
    #
    #         fname = f'{cropped_dir}_dila/cropped_image_{row[0]}.tif'
    #         if os.path.isfile(fname):
    #             # pixel_num = main_point.get_width(fname, 3, 6, 30, 2)
    #             poly_num = len(row[3].split(','))
    #             distance = int(row[2])
    #             if poly_num > 4 or distance > 250:
    #                 pixel_num = 999
    #                 row.append(str(pixel_num))  # 将 pixel_num 添加到行末
    #             else:
    #                 pixel_num = inner_circle.inner_circle(fname, circle_dir)
    #                 row.append(str(pixel_num))  # 将 pixel_num 添加到行末
    #         else:
    #             print('路径不存在：' + fname)
    #         # pixel_num = main_point.get_width(f'luopu/cropped/cropped_image_{row[0]}.tif', 3, 6, 30, 2)
    #         # row.append(str(pixel_num))  # 将 pixel_num 添加到行末
    #         # red_point.get_red(f'luopu/cropped/cropped_image_{row[0]}.tif')
    #         # get_direction.get_direction(f'cropped/cropped_image_{row[0]}.tif')
    # with open(output_dir, 'w', newline='', encoding='gbk') as csvfile:
    #     writer = csv.writer(csvfile)
    #     writer.writerows(rows)


csv_file = 'csvs/xn_ded.csv'  # csv路网的路径
tif_file = 'tifs/xn_pan_已修改.tif'  # tif预测结果的路径
output_dir = 'csvs/xinnan_gbk_circle_dila.csv'  # 输出路宽信息的路径
cropped_dir = 'ded/cropped_pan'  # 根据路网切割的TIF路径
circle_dir = 'ded/circle'  # 内接圆图像路径
crop_by_csv(csv_file, tif_file, cropped_dir, output_dir, circle_dir)
